import os
from multiprocessing import RLock
import numpy as np
import cv2
from  datetime import datetime as dt


class ImagesManager:
    """
    ImageManager
    """
    def __init__(self, root, name_pattern=None):
        """
        :param root 保存图片的根目录
        :param name_pattern  文件名的模式，可以使用year、month、day、hour, minute, second, seq和info等模式变量
        """
        self.root = to_directory(root)
        self.name_pattern = '{seq:07d}.png' if name_pattern is None else name_pattern

        self.lock = RLock()
        self.current_dir = None
        self.current_seq = None

        # self.year, self.month, self.day, self.hour, self.minute, self.second = get_date()

    def save(self, img: np.ndarray, datetime=None, info:str=None):
        """
        把文件保存在图片根目录下<year>/<year>-<month>/<year>-<month>-<day>/子目录中，
        该子目录自动创建，图片依次被保存为：
        0000001.jpg
        0000002.jpg
        0000003.jpg
        ......
        其中<month>和<day>从1开始数。
        :param img 被保存的图片，类型是np的矩阵
        :param datetime 图片的时间，datetime.datetime类型。缺省值None表示当前时间
        :param info 被保存图片的辅助信息，如作者、创作日期等。None表示没有辅助信息
        :return 返回图片的id，图片id是形如<year>-<month>-<day>-0000003的字符串，注意，id不等于文件名，不含有扩展名
        """
        if info is None:
            info = ''
        else:
            info = ''.join(['_' if c in r'\/:*?"|' else c for c in info])

        if datetime is None:
            datetime = dt.now()
        with self.lock:
            path = self.get_directory(datetime.year, datetime.month, datetime.day)
            seq = self._next_seq(path)

        pic_name = self.name_pattern.format(year=datetime.year, month=datetime.month, day=datetime.day, hour=datetime.hour,
                                            minute=datetime.minute, second=datetime.second, info=info, seq=seq)
        path = path + os.sep + pic_name
        # 保存图片
        cv2.imwrite(path, img)
        id = self._get_id(datetime, pic_name)
        return id

    def _next_seq(self, directory):
        if self.current_dir is None or self.current_dir != directory:
            os.makedirs(directory, exist_ok=True)
            self.current_seq = len(os.listdir(directory))
        self.current_seq += 1
        return self.current_seq

    def _get_id(self, datetime, name):
        return "%s-%02d-%02d/%s" % (datetime.year, datetime.month, datetime.day, name)

    def load(self, image_id):
        """
        根据image_id返回对应文件中保存的图片
        """
        full_path = self.get_full_path(image_id)
        img = cv2.imread(full_path)
        return img

    def remove(self, image_id):
        """
        删除指定 image_id 的图片
        """
        path = self.get_full_path(image_id)
        if not os.path.exists(path):
            return False
        os.remove(path)
        return True

    def get_full_path(self, image_id):
        year = int(image_id[0:4])
        month = int(image_id[5:7])
        day = int(image_id[8:10])
        name = image_id[11:]
        return self.get_directory(year, month, day) + os.sep + name

    def get_directory(self, year, month, day):
        year = str(year)
        month = '%02d' % month
        day = '%02d' % day

        first_level = year
        second_level = first_level + '-' + month
        third_level = second_level + '-' + day
        path = self.root + first_level + os.sep + second_level + os.sep + third_level
        return path


def to_directory(directory):
    d = directory.strip()
    if d[-1] not in ('\\', '/'):
        d += os.sep
    return d


def exists(directory):
    d = to_directory(directory)
    return os.path.exists(d)


if __name__ == '__main__':
    im = ImagesManager('../../temp', '{info}-{seq:07d}.jpg')
    # img = np.random.randint(0, 255, [300, 200, 3])
    # id = im.save(img, 'sub', 'ttt')
    # print(id)

    # img = im.load('2022-06-20/sub/ttt-0000002.jpg')
    # cv2.imshow('aaa', img)
    # cv2.waitKey()
    # cv2.destroyAllWindows()

    im.remove('2022-06-20/sub/ttt-0000003.jpg')